import torch
from torch.optim import Optimizer
import sys

class GradientDescent(Optimizer):
    def __init__(self, params, lr=0.01):
        if lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        defaults = dict(lr=lr)
        super(GradientDescent, self).__init__(params, defaults)
    def set_lr(self, new_lr):
        for group in self.param_groups:
            group['lr'] = new_lr

    def step(self, closure=None):
            if closure is not None:
                with torch.enable_grad():
                    loss = closure()
            for group in self.param_groups:
                for p in group['params']:
                    if p.grad is None:
                        continue
                    with torch.no_grad():
                        print("p.data",p.data)
                        print("group['lr']",group['lr'])
                        print("p.grad",p.grad)

                        p.data -= group['lr'] * p.grad
                        

